#!/usr/bin/env python2

#==============================================================#
# File      :   pg-register
# Ctime     :   2020-05-16
# Mtime     :   2020-12-17
# Desc      :   register or adjust pgsql consul service
# Path      :   /pg/bin/pg-register
# Depend    :   pg-role
# Author    :   Vonng(fengruohang@outlook.com)
# Copyright (C) 2018-2022 Ruohang Feng
#==============================================================#


"""
keep dcs service definition role sync with real database role
"""

import os, sys, json


def dict_compare(d1, d2):
    for k in d1.keys():
        if not d2.has_key(k):
            return False
        else:
            if type(d1[k]) is dict:
                dict_compare(d1[k], d2[k])
            else:
                if d1[k] != d2[k]:
                    return False
    return True


def modify_value(s, role):
    if s == 'primary' or s == 'replica':
        return role

    # common service naming pattern: <cluster>-primary <cluster>-replica
    if s.endswith('primary') or s.endswith('replica'):
        return s.rstrip('primary').rstrip('replica') + role

    # other case should be use with cautious, what if service name contains primary and replica ?
    if 'primary' in s or 'replica' in s:
        return s \
            .replace("primary", "[NEW_ROLE_PLACEHOLDER]") \
            .replace("replica", "[NEW_ROLE_PLACEHOLDER]") \
            .replace("[NEW_ROLE_PLACEHOLDER]", role)
    return s


def modify_role(svc, role):
    changed = False
    if "service" not in svc:  # not a valid service
        return svc

    svc_name, svc_tags, svc_meta = svc["service"].get("name"), svc["service"].get("tags"), svc["service"].get("meta")
    if not svc_name: return svc

    if svc_tags:
        new_tags = []
        for tag in svc_tags:
            new_tag = modify_value(tag, role)
            new_tags.append(new_tag)
            if new_tag != tag:
                changed = True
                print("[item-changed] change svc-%s tags %s to %s" % (svc_name, tag, new_tag))
        svc["service"]["tags"] = new_tags

    if svc_meta:
        new_meta = {}
        for k, v in svc_meta.items():
            new_v = modify_value(v, role)
            new_meta[k] = new_v
            if new_v != v:
                changed = True
                print("[item-changed] change svc-%s meta %s from %s to %s" % (svc_name, k, v, new_v))
        svc["service"]["meta"] = new_meta

    return json.loads(json.dumps(svc)), changed


def anti_entropy(designate_role=None, conf_dir="/etc/consul.d"):
    # process consul service definition
    files = os.listdir(conf_dir)
    for fname in files:
        if not fname.startswith('svc') or not fname.endswith('.json'):
            continue
        filepath = os.path.join(conf_dir, fname)  # load svc-*.json
        svc = json.load(open(filepath))  # load json with top level key "service"
        new_svc, changed = modify_role(svc, designate_role)  # modify role related attribte

        if changed:  # if any change needs to be made
            print("[svc-changed]: %s" % filepath)
            with open(filepath, 'w') as dst:  # and dump it back
                json.dump(new_svc, dst, indent=4)
        else:
            print("[svc-intact]: %s" % filepath)


if __name__ == '__main__':
    role = None
    conf_d = "/etc/consul.d"
    if len(sys.argv) > 1:
        if sys.argv[1] in ['primary', 'replica']:
            role = sys.argv[1]
        elif sys.argv[1] in ['unknown']:
            print("role is unknown, do nothing")
            exit(0)
        else:
            print("allowed roles: primary, replica")
            exit(1)
    else:
        print("role is required")
        exit(2)

    if len(sys.argv) > 2 and sys.argv[2] != "":
        conf_d = sys.argv[2]

    anti_entropy(role, conf_d)
